# 文件    ： ai-web.py
# 时间    ： 2025/3/21 15:47
# 作者    ： Luzhaohui
# 环境    ： PyCharm
import time

import uvicorn
import contextlib
import threading
from pydantic import BaseModel
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from transformers import Qwen3VLForConditionalGeneration, Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
import torch

# WEB 服务
from webUI import Ui_MainWindow
from conf import *
from file import *
from screen import *
from dt import *
import sys, os
import base64
from PIL import Image
from io import BytesIO

VER = '1.5.11.19'


class UvicornServer(uvicorn.Server):
    def install_signal_handlers(self):
        pass

    # 上下文管理
    @contextlib.contextmanager
    def run_in_thread(self):
        print('uvicorn_server start thread 1')
        thread = threading.Thread(target=self.run)
        thread.start()
        try:
            while not self.started:
                time.sleep(1e-3)
            yield
        finally:
            self.should_exit = True
            thread.join()
        print('uvicorn_server start thread 2')


global_model = None
global_processor = None


class WebApi(QThread):
    update_item = pyqtSignal(list)
    chat_url = ''

    def __init__(self, model_path, model_name, chat_url, web_host, web_port, max_tokens, max_width, max_height):
        super().__init__()
        self.model_path = model_path
        self.model_name = model_name
        self.chat_url = chat_url
        self.web_host = web_host
        self.web_port = web_port
        self.max_tokens = max_tokens
        self.max_width = max_width
        self.max_height = max_height
        self.stop = False

    def is_second(self, second):
        if 'result' in second.keys() and 'text' in second.keys() and 'default' in second.keys():
            if len(second['result']) >= 2 and len(second['text']) >= 2 and second['default'] != '':
                return True
        return False

    def chat_reason(self, image, text, processor, model):
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": image,
                    },
                    {
                        "type": "text",
                        "text": text
                    },
                ],
            }
        ]
        # Preparation for inference
        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        image_inputs, video_inputs = process_vision_info(messages)
        inputs = processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        inputs = inputs.to(model.device)

        # Inference: Generation of the output
        generated_ids = model.generate(**inputs, max_new_tokens=self.max_tokens)

        generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
        output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True,
                                             clean_up_tokenization_spaces=False)
        return output_text

    def run(self):
        model_dir = AIPATH + '%s/%s' % (self.model_path, self.model_name)
        model_name = os.path.basename(self.model_name).lower()
        is_qwen3 = 'qwen3-vl' == model_name[:8]
        print("start Model/Processor %s %s" % (model_name, str(is_qwen3)))
        global global_model
        global global_processor
        if global_model is None:
            print('load model')
            if is_qwen3:
                global_model = Qwen3VLForConditionalGeneration.from_pretrained(
                    model_dir,
                    dtype=torch.bfloat16,
                    device_map="auto",
                )
            else:
                global_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
                    model_dir,
                    dtype=torch.bfloat16,
                    device_map="auto",
                )
        if global_processor is None:
            print('auto processor %d %d' % (self.max_width, self.max_height))
            global_processor = AutoProcessor.from_pretrained(model_dir, max_pixels=self.max_width * self.max_height)

        app = FastAPI()

        class Item(BaseModel):
            text: str
            second: dict
            image: str

        @app.post("/model")
        async def get_model():
            data = {
                "object": "list",
                "data": [
                    {
                        "id": self.model_name,
                    }
                ]
            }
            return JSONResponse(content=data)

        @app.post("/generate-text")
        async def chat(item: Item):
            """
            Endpoint to handle user input (text + image) and generate text results.
            """
            image_data = item.image.split(',')
            if len(image_data) >= 2 and 'base64' in image_data[0]:
                self.update_item.emit([item.text, image_data[0]])
            else:
                self.update_item.emit([item.text, item.image])
            start_time = time.perf_counter()

            global global_model
            global global_processor

            chat_second = False
            while True:
                # print(item.second)
                output_title = ''
                output_top = ''
                # first chat
                output_text = self.chat_reason(item.image, item.text, global_processor, global_model)
                print(output_text)
                if not self.is_second(item.second):
                    break
                # second 解析第一次标题
                firstResult = eval(output_text[0].lstrip('```json').rstrip('```'))
                if type(firstResult) != dict:
                    break
                for k, v in firstResult.items():
                    if '标题' == k or 'title' == k:
                        output_title = str(v)
                        continue
                    if '抬头' in k:
                        output_top = str(v)
                        continue
                print('second title = %s， top = %s' % (output_title, output_top))
                if output_title == '':
                    break
                resultList = item.second['result']
                textList = item.second['text']
                textDefault = item.second['default']
                second_index = -1
                for index, result in enumerate(resultList):
                    # print(result)
                    if result == '':
                        # 空
                        break
                    # 判断结果,多个条件
                    result_slist = result.split(';')
                    if result_slist[0] in output_title:
                        if len(result_slist) == 1:
                            # 1匹配
                            second_index = index
                            break
                        if len(result_slist) == 2 and output_top != '':
                            # 2匹配
                            if result_slist[1] in str(output_top):
                                second_index = index
                                break
                if second_index == -1:
                    second_text = textDefault
                else:
                    second_text = textList[second_index]
                print('second text = %s' % second_text)
                output_text = self.chat_reason(item.image, second_text, global_processor, global_model)
                chat_second = True
                print(output_text)
                break

            proc_time = time.perf_counter() - start_time
            print(proc_time)

            proc_id = now_datetime_int()
            total_tokens = 0
            self.update_item.emit([proc_id, proc_time, total_tokens])
            if output_title == '':
                return JSONResponse(content={"output_text": output_text, "chat_type": chat_second})
            else:
                return JSONResponse(
                    content={"title": output_title, "output_text": output_text, "chat_type": chat_second})

        # 阻塞启动
        # uvicorn.run(self.app, host=self.web_host, port=self.web_port)
        # 非阻塞启动
        config = uvicorn.Config(app, host=self.web_host, port=self.web_port)
        server = UvicornServer(config=config)
        with server.run_in_thread():
            self.update_item.emit(['start'])
            while True:
                if self.stop:
                    break
                QThread.msleep(10)
            self.update_item.emit(['exit'])

    def exit(self):
        self.stop = True


class WebRun(QMainWindow, Ui_MainWindow):
    # 窗口
    stdoutBak = ''
    stderrBak = ''
    imagePath = ''
    resultPath = ''
    # 模型
    modelPath = ''
    modelName = ''
    systemPrompt = ''
    # 参数
    stream = False
    max_tokens = 1024
    max_width = 1280
    max_height = 784
    # WEB
    isWeb = False
    webHost = '127.0.0.1'
    webPort = 8099
    webChatUrl = '/aiChat'
    webModelUrl = '/model'
    webThread = None
    # 数据库
    isMysql = False
    mysqlHost = '127.0.0.1'
    mysqlPort = 3306
    mysqlUser = 'root'
    mysqlPass = 'root'
    mysqlDatabase = 'LLM'
    msql = None
    start = False
    # 技术
    queue_nums = 0
    reason_nums = 0

    def __init__(self, parent=None):
        super(WebRun, self).__init__(parent)
        save_log('start')
        self.setupUi(self)
        self.init_redirection()
        self.config = WebConf('web.ini')
        self.msql = None
        self.web = None
        self.init_var()
        self.init()
        self.init_ui()
        self.init_mysql()
        self.init_web()
        if self.config.varConf['IsStart']:
            self.start_web()

    def init_redirection(self):
        self.stdoutBak = sys.stdout
        self.stderrBak = sys.stderr

    def reset_redirection(self):
        sys.stdout = self.stdoutBak
        sys.stderr = self.stderrBak

    def closeEvent(self, event, **kwargs):
        if not self.start and showBoxYes(self, "系统将退出，是否确认?  "):
            self.exit()
            event.accept()
        else:
            event.ignore()

    def exit(self):
        save_log('exit')
        if self.msql is not None:
            del self.msql
        if self.web is not None:
            del self.web
        del self.config
        self.reset_redirection()

    def init_var(self):
        # 路径
        self.imagePath = self.config.varConf['ImagePath']
        self.resultPath = self.config.varConf['ResultPath']
        # 模型
        self.modelPath = self.config.varModel['ModelPath']
        self.modelName = self.config.varModel['ModelName']
        self.systemPrompt = self.config.varModel['SystemPrompt']
        # 参数
        self.stream = self.config.varPara['stream']
        self.max_tokens = self.config.varPara['max_tokens']
        self.max_width = self.config.varPara['max_width']
        self.max_height = self.config.varPara['max_height']
        # WEB
        self.isWeb = self.config.varWeb['IsWebServer']
        self.webHost = self.config.varWeb['WebHost']
        self.webPort = self.config.varWeb['WebPort']
        self.webChatUrl = self.config.varWeb['WebChatUrl']
        self.webModelUrl = self.config.varWeb['WebModelUrl']
        # 数据库
        self.isMysql = self.config.varMysql['IsMysql']
        self.mysqlHost = self.config.varMysql['MysqlHost']
        self.mysqlPort = self.config.varMysql['MysqlPort']
        self.mysqlUser = self.config.varMysql['MysqlUser']
        self.mysqlPass = self.config.varMysql['MysqlPass']
        self.mysqlDatabase = self.config.varMysql['MysqlDatabase']
        self.msql = None
        #
        self.start = False

    def init(self):
        self.setWindowIcon(QIcon(APPLOGO))
        self.setWindowTitle(' %s ver[%s]' % (self.config.varConf['Title'], VER))
        # 路径
        self.lineEditModel.setText(self.modelName)
        self.lineEditImageDir.setText(self.imagePath)
        # WEB
        self.checkBoxWebServer.setChecked(self.isWeb)
        self.spinBoxWebPort.setValue(self.webPort)
        # 数据
        self.checkBoxMysql.setChecked(self.isMysql)
        self.lineEditMysqlInfo.setText(
            '%s@%s:%d->%s' % (self.mysqlUser,
                              self.mysqlHost,
                              self.mysqlPort,
                              self.mysqlDatabase))
        # 定时器
        # self.timer = QTimer()
        # self.timer.timeout.connect(self.update_time)
        # self.timer.start(1000)

    def init_ui(self):
        # 选择路径
        self.pushButtonSelectDir.clicked.connect(self.change_data_dir)
        # WEB
        self.checkBoxWebServer.clicked.connect(self.change_web)
        # 数据库
        self.checkBoxMysql.clicked.connect(self.change_mysql)
        # 启动
        self.pushButtonReasonService.clicked.connect(self.start_web)
        self.flush_show()

    def init_mysql(self):
        if self.msql is not None:
            return

    def init_web(self):
        if self.web is not None:
            return
        # self.web = WebApi('%s/%s' % (self.modelPath, self.modelName),
        #                   self.webChatUrl,
        #                   self.webHost,
        #                   self.webPort)

    def change_data_dir(self):
        # 选择数据目录
        dir_name = AIPATH + self.imagePath
        path = QFileDialog.getExistingDirectory(self, " 选择图片主目录 ", dir_name)
        if os.path.exists(path):
            self.imagePath = path[len(AIPATH):]
            self.config.write('CONF', 'ImagePath', self.imagePath)
            self.lineEditImageDir.setText(self.imagePath)
            main_dir = os.path.dirname(self.imagePath)
            path = main_dir + '/' + os.path.basename(self.resultPath)
            self.config.write('CONF', 'ResultPath', path)

    def change_web(self):
        if self.isWeb != self.checkBoxWebServer.isChecked():
            self.isWeb = self.checkBoxWebServer.isChecked()
            self.config.write('WEB', 'IsWebServer', self.isWeb)

    def change_mysql(self):
        if self.isMysql != self.checkBoxMysql.isChecked():
            self.isMysql = self.checkBoxMysql.isChecked()
            self.config.write('MYSQL', 'IsMysql', self.isMysql)

    def flush_show(self):
        self.pushButtonSelectDir.setEnabled(not self.start)
        self.checkBoxWebServer.setEnabled(not self.start)
        self.checkBoxMysql.setEnabled(not self.start)

    def start_web(self):
        self.start = not self.start
        if self.start:
            # 启动
            self.pushButtonReasonService.setText('推理服务关闭')
            self.web = WebApi(self.modelPath,
                              self.modelName,
                              self.webChatUrl,
                              self.webHost,
                              self.webPort,
                              self.max_tokens,
                              self.max_width,
                              self.max_height)
            self.web.update_item.connect(self.show_web)
            self.web.start()
        else:
            self.pushButtonReasonService.setText('推理服务启动')
            self.web.exit()
        self.flush_show()

    def show_web(self, info):
        # web线程状态
        # print(info)
        if len(info) == 1:
            self.labelReasonStatus.setText('')
            if info[0] == 'start':
                self.pushButtonReasonService.setText('推理服务关闭')
            if info[0] == 'exit':
                self.pushButtonReasonService.setText('推理服务启动')
        elif len(info) == 2:
            # 请求
            self.lineEditReasonText.setText(info[0])
            self.lineEditReasonImage.setText(info[1])
            self.labelReasonStatus.setText('推理中...')
            self.queue_nums += 1
            self.labelQueueNums.setText('%d' % self.queue_nums)
        elif len(info) == 3:
            # 结束
            self.queue_nums -= 1
            self.labelQueueNums.setText('%d' % self.queue_nums)
            self.reason_nums += 1
            self.labelReasonNums.setText('%d' % self.reason_nums)
            self.labelReasonStatus.setText('%d %.3f %d' % (info[0], info[1], info[2]))


if __name__ == '__main__':
    print('web service')
    app = QApplication(sys.argv)

    web = WebRun()
    web.show()

    sys.exit(app.exec_())
